import torch
import torch.nn as nn
import torchvision

# hyper-param
cls_num = 751  # market1501 train 751

# model
# net = torchvision.models.resnet18(pretrained=True)
net = torchvision.models.resnet50(pretrained=True)
# net = torchvision.models.resnet101(pretrained=True)
# net = torchvision.models.resnet152(pretrained=True)
in_feat = net.fc.in_features
new_fc = nn.Linear(in_features=in_feat, out_features=cls_num)
net.fc = new_fc

# use demo
if __name__ == "__main__":
    device = torch.device("cpu")  # 需要改一下
    net.to(device)
    x = torch.ones((2, 3, 384, 128), dtype=torch.float32).to(
        device)  # 必须接收3通道图
    y = net(x)  # 前向传播
    print(y.shape)  # [batch_size, cls_num]
